// ==================================================================================
// Shared Genomics Project MPI Codebase
// Version 1.0 30/04/2010
//
// (c) 2010 University of Manchester all rights reserved
//
// This file is distributed under the GNU General Public License, Version 2.  
// Please see the file COPYING.txt for more details
// ==================================================================================

/*!
\file 
\brief Parrellised Filtering Functions 
*/

#include <string.h>
#include <time.h>

#include "bfilter.h"
#include "copyfile.h"
#include "bmapfile.h"
#include "bpedfile.h"
#include "pmaths.h"
#include "suffixes.h"

int bfilter_parse_exclude_list(struct selected_options *ops, int nSamples, BOOL *persons_to_del) {
	FILE *in = NULL;
	char line[80];

	if (ops == NULL || nSamples <= 0 || persons_to_del == NULL) return 0;
	else if (strlen(ops->szExcludeListFileName) == 0) return 0;

	in = fopen(ops->szExcludeListFileName, "r");
	if (in != NULL) {
		while (fgets(line, 80, in)) {
			int person_index = atoi(line);
			if (person_index >= 0 && person_index <= nSamples) persons_to_del[person_index] = TRUE;
		}
		fclose(in);
	}

	return 1; // N/B Optional file so no need to flag if failed.
}

int bfilter_person_do_mind(int nSamples, double *missingness, BOOL *persons_to_del, struct selected_options *ops) {
	int i;

	if (nSamples <= 0 || missingness == NULL || persons_to_del == NULL || ops == NULL)  return 0;

	if (ops->MAX_IND_MISSING < 1.0) {
		for (i = 0; i < nSamples; i++) if (missingness[i] > ops->MAX_IND_MISSING) persons_to_del[i] = TRUE;
	}
	
	return 1;
}

int bfilter_snp_do_maf_and_mind(char *bpedData, struct small_locus *snps, int nSNPs, BOOL *snps_to_del, char *snps_reasons_to_del, double *snp_del_stats, int nSamples, BOOL *persons_to_del, BOOL *sex, BOOL *founders, struct chromosomes *chromos, struct selected_options *ops) {
	int i, j, nActualSamples = 0;
	char *pBpedData = NULL;
	int maf_count = 0;

	if (bpedData == NULL || snps == NULL || nSNPs <= 0 || snps_to_del == NULL || snps_reasons_to_del == NULL || snp_del_stats == NULL ||
		nSamples <= 0 || persons_to_del == NULL || sex == NULL || founders == NULL ||
		chromos == NULL || ops == NULL) return 0;
	else if (!chromos->isdefined) return 0;

	// Count the actual number of samples, i.e. people flagged out to be deleted.
	// If zero, flag all the SNPs for deletion as no more information in the input data-set.
	for (i = 0; i < nSamples; i++) if (!persons_to_del[i]) nActualSamples++;
	if (nActualSamples == 0) {
		for (i = 0; i < nSNPs; i++) {
			snps_to_del[i] = TRUE;
			snps_reasons_to_del[i] = FILTER_REASON_PERSON_MISSINGNESS;
		}
		return 1;
	}
	
	// Iterate over each individual in the sample
	for (i = 0, pBpedData = bpedData; i < nSNPs; i++) {
		int geno_nm = 0; // count 1 per genotype, for missingness
		int geno_real = ops->oblig_missing ? 0 : nActualSamples; // count 1 per non-obligatory missing genotype (or set to N individuals)
		double snp_genotyping = 0.0;
		BOOL X = FALSE, haploid = FALSE;
		int nm = 0; // count 1 per allele, for frequency
		struct small_locus *loc = &snps[i];
		double bp = 0, pos = 0, freq = 0;

		// Already removed by some other filter so skip to next SNP record.
		if (snps_to_del[i]) {
			pBpedData += nSamples * 2 * sizeof(char);
			continue; 
		}

		// Determine type of SNP
		if (chromos->chr_sex[loc->chr]) X = TRUE;
		else if (chromos->chr_haploid[loc->chr]) haploid = TRUE;

		// Iterate over each individual in the sample
		for (j = 0; j < nSamples; j++, pBpedData += (2 * sizeof(char))) {
			BOOL s1, s2;

			if (persons_to_del[j]) continue;
			s1 = pBpedData[0];
			s2 = pBpedData[1];

			// Check female Y genotypes
			if (chromos->chr_Y[loc->chr] && !sex[j]) {
				// Set to missing, unless in a RECODE mode.
				if (!ops->preserve_all_genotypes) {
					s1 = pBpedData[0] = TRUE;
					s2 = pBpedData[1] = FALSE;
					continue; // Skip to next person.
				}
			}
			
			// For haploid heterozygosity check, also consider all individuals
			if (haploid || (X && sex[j])) {
				if (!s1 && s2) {
					// Set to missing unless in RECODE mode
					if (!ops->preserve_all_genotypes) {
						pBpedData[0] = TRUE;
						pBpedData[1] = FALSE;
					}
				}
			}
			
			// For missing genotypes
			if (!(s1 && !s2)) geno_nm++;

			// For allele frequencies only consider founders.
			if (ops->summ_nonfounders || founders[j]) {
				if (haploid || (X && sex[j])) {
					// Haploid counts
					// "1" allele 
					if (!s1 && !s2) {  // FF = hom(11)
						freq++;
						nm++;
					} else if (s1 && s2) { // TT = hom(22)
						nm++;
					}
				} else {
					// Autosomal count
					// "1" allele 
					if (!s1) { 
						if(!s2) {	// 00 = hom(11)
							freq += 2;
							nm += 2;
						} else {  // 01 = het(12)
							freq += 1;
							nm += 2;
						}
					} else if (s2) {
						nm += 2;
					}
				}
			} 
		}
		
		// Calculate allele frequencies
		if (ops->af_count) { // Allele counts...
			bp = (long int) (nm - freq);            // Use freq to store count (keep as is)	
			pos = geno_real - geno_nm; // Use "pos" to store number missing genotypes
		} else {
			// ... or frequencies
			if (nm > 0) freq /= nm; 
			else freq = 1;
		}

		// Record total proportion of missingness.
		snp_genotyping =  nActualSamples > 0 ? (double) geno_nm/ (double) geno_real : 0;
		// Exclude if SNP has too many missing genotypes
		if (snp_genotyping < (1- ops->MAX_GENO_MISSING)) {
			snps_to_del[i] = TRUE;
			snps_reasons_to_del[i] = FILTER_REASON_SNP_MISSINGNESS;
			snp_del_stats[i] = 1 - snp_genotyping;
		}

		// Make allele1 always the least common allele
		if (!ops->af_count && freq > 0.5) {
			char tmp[ALLELE_LENGTH];
			char *pBpedRecord = bpedData + (nSamples * 2 * sizeof(char) * i); // Go to the current SNP record & adjust the ordering.

			// then we need to swap alleles
			freq = 1 - freq;
			strcpy(tmp, loc->allele2);
			strcpy(loc->allele2, loc->allele1);
			strcpy(loc->allele1, tmp);

			// Update the data in the BPED array.
			for (j = 0; j < nSamples; j++, pBpedRecord += (2 * sizeof(char))) {
				if (persons_to_del[j]) continue;
				if (pBpedRecord[0] == pBpedRecord[1]) {
					pBpedRecord[0] = !pBpedRecord[0];
					pBpedRecord[1] = !pBpedRecord[1];
				}
			}	
		} 
		
		// Remove rare SNPs
		// Note epsilon correction for MAF, due to floating point 
		// issues: only apply to the lower MAF range
		if (freq < 0 || (freq + ops->epsilon) < ops->min_af || freq > ops->max_af) {
			snps_to_del[i] = TRUE;
			snps_reasons_to_del[i] = FILTER_REASON_SNP_MAF;
			snp_del_stats[i] = freq;
			maf_count++;
		}
	}
	
	return 1;
}

int bfilter_snp_do_hwe(char *bpedData, struct small_locus *snps, int nSNPs, BOOL *snps_to_del, char *snps_reasons_to_del, double *snp_del_stats, int nSamples, BOOL *persons_to_del, BOOL *sex, BOOL *founders, double *phe_values, struct chromosomes *chromos, struct selected_options *ops) {
	char *pBpedData = NULL;
	int i, j;
	
	if (bpedData == NULL || snps == NULL || nSNPs <= 0 || snps_to_del == NULL || snps_reasons_to_del == NULL || snp_del_stats == NULL ||
		nSamples <= 0 || persons_to_del <= 0 || sex == NULL || founders == NULL ||
		chromos == NULL || ops == NULL || phe_values == NULL) return 0;
	else if (!chromos->isdefined) return 0;
	
	// Loo through the SNPs
	for (i = 0, pBpedData = bpedData; i < nSNPs; i++) {
		struct small_locus *loc = &snps[i];
		
		// Compute p-values for HWE test in cases, controls & all
		// Only consider founders
		int a11 = 0, a12 = 0, a22 = 0;
		int u11 = 0, u12 = 0, u22 = 0;
		int b11 = 0, b12 = 0, b22 = 0;

		// Allele frequencies
		double afreq = 0, ufreq = 0, freq = 0;
		BOOL include_cases = TRUE, include_controls = TRUE;

		BOOL X = FALSE, haploid = FALSE;

		// Already removed by some other filter so skip to next SNP record.
		if (snps_to_del[i]) {
			pBpedData += nSamples * 2 * sizeof(char);
			continue; 
		}

		if (chromos->chr_sex[loc->chr]) X = TRUE;
		else if (chromos->chr_haploid[loc->chr]) haploid = TRUE;
		
		// Iterate over each individual, founders only
		for (j = 0; j < nSamples; j++, pBpedData += (2 * sizeof(char))) {
			if (persons_to_del[j]) continue;
			
			// Only consider founders, & diploid genotypes
			if (ops->summ_nonfounders || founders[j]) {
				if ( !(haploid || (X && sex[j]))) {	
					BOOL s1 = pBpedData[0], s2 = pBpedData[1];
					
					// Consider everybody, irrespective of phenotype
					// (QT, C/C or missing)
					if (!s1) { 
						if (!s2) b11++;   //   00 = hom(11)
						else b12++;       //   01 = het(12)
					}
					else if ( s2 ) b22++; // 11 = hom(22)

					// for binary trait, separately for cases/controls
					if (ops->bt)  {
						if (phe_values[j] == 1) {
							if (!s1) { 
								if (!s2) u11++;   //   00 = hom(11)
								else u12++;       //   01 = het(12)
							}
							else if (s2) u22++; //   11 = hom(22)

						}
						else if (phe_values[j] == 2) {
							if (!s1) { 
								if (!s2) a11++;   //   00 = hom(11)
								else a12++;         //   01 = het(12)
							}
							else if ( s2 ) a22++; //   11 = hom(22)
						}
					}
				}
			}
		}
		
		// Allele frequencies
		if (ops->qt)
			freq = ( b11 + (double)b12/2.0 ) / (double) (b11 + b12 + b22);
		else {
			afreq = (a11 + (double) a12/2.0) / (double) (a11 + a12 + a22);
			ufreq = (u11 + (double) u12/2.0) / (double) (u11 + u12 + u22);
			freq =  (b11 + (double) b12/2.0) / (double) (b11 + b12 + b22);

			if (a11 + a12 + a22 == 0 ) include_cases = FALSE;
			if (u11 + u12 + u22 == 0 ) include_controls = FALSE;
		}

		if (ops->qt) {
			double p;

			if (ops->HWD_standard) {
				double tot = b11 + b12 + b22;
				double exp_11 = freq * freq * tot;
				double exp_12 = 2 * freq * (1-freq) * tot;
				double exp_22 = (1-freq) * (1-freq) * tot;

				double chisq = ( (b11-exp_11)*(b11-exp_11) ) / exp_11 
					+ ( (b12-exp_12)*(b12-exp_12) ) / exp_12 
					+ ( (b22-exp_22)*(b22-exp_22) ) / exp_22 ;

				p = chiprobP(chisq,1);
			}
			else
				p = SNPHWE( b12, b11, b22 );

			if (p <= ops->HWD_limit ) {
				snps_to_del[i] = TRUE; // Set to delete SNP
				snps_reasons_to_del[i] = FILTER_REASON_SNP_HWE;
				snp_del_stats[i] = p;
			}
		} 
		else
		{
			// For case/control data
			double p, p_a, p_u;

			if (ops->HWD_standard) {
				double exp_a11 = afreq * afreq * (a11+a12+a22);
				double exp_a12 = 2 * afreq * (1-afreq) * (a11+a12+a22);
				double exp_a22 = (1-afreq) * (1-afreq) * (a11+a12+a22);

				double exp_u11 = ufreq * ufreq * (u11+u12+u22);
				double exp_u12 = 2 * ufreq * (1-ufreq) * (u11+u12+u22);
				double exp_u22 = (1-ufreq) * (1-ufreq) * (u11+u12+u22);

				double exp_11 = freq * freq * (b11+b12+b22);
				double exp_12 = 2 * freq * (1-freq) * (b11+b12+b22);
				double exp_22 = (1-freq) * (1-freq) * (b11+b12+b22);

				double chisq_a = ( (a11-exp_a11)*(a11-exp_a11) ) / exp_a11 
					+ ( (a12-exp_a12)*(a12-exp_a12) ) / exp_a12 
					+ ( (a22-exp_a22)*(a22-exp_a22) ) / exp_a22 ;

				double chisq_u = ( (u11-exp_u11)*(u11-exp_u11) ) / exp_u11 
					+ ( (u12-exp_u12)*(u12-exp_u12) ) / exp_u12 
					+ ( (u22-exp_u22)*(u22-exp_u22) ) / exp_u22 ;

				double chisq = ( (b11-exp_11)*(b11-exp_11) ) / exp_11 
					+ ( (b12-exp_12)*(b12-exp_12) ) / exp_12 
					+ ( (b22-exp_22)*(b22-exp_22) ) / exp_22 ;

				p = chiprobP(chisq,1);
				p_a = chiprobP(chisq_a,1);
				p_u = chiprobP(chisq_u,1);
			} else {
				p = SNPHWE( b12, b11, b22 );
				p_a = SNPHWE( a12, a11, a22 );
				p_u = SNPHWE( u12, u11, u22 );
			}

			// Controls (and, if possible, exclude on this value)
			if (include_controls && p_u < ops->HWD_limit && p > -1) {
				if (!ops->HWD_filter_on_all) { 
					snps_to_del[i] = TRUE;
					snps_reasons_to_del[i] = FILTER_REASON_SNP_HWE;
					snp_del_stats[i] = p_u;
				}
			}

			// In total sample, and if needed, exclude here
			if (p < ops->HWD_limit && p > -1) { 
				if (ops->HWD_filter_on_all || !include_controls) {
					snps_to_del[i] = TRUE;
					snps_reasons_to_del[i] = FILTER_REASON_SNP_HWE;
					snp_del_stats[i] = p;
				}
			}
		}
	}

	return 1;
}

// Only log for a filter which has been run.
int bfilter_write_log(struct selected_options *ops) {
	FILE *out = NULL;
	char localFilepath[FILE_PATH_LENGTH], remoteFilepath[FILE_PATH_LENGTH];
	time_t now = time(NULL);

	if (ops == NULL) return 0;
	
	sprintf(localFilepath, "%s%s%s", ops->ldir, ops->szRunId, SUFFIX_FILTER_SETTINGS);
	sprintf(remoteFilepath, "%s%s%s", ops->rdir, ops->szRunId, SUFFIX_FILTER_SETTINGS);
	out = fopen(localFilepath, "w");
	if (out == NULL) return 0;

	// Record the options & data set involved when filter was invoked.
	fprintf(out, "Dated=%s", ctime(&now));
	fprintf(out, "id=%s\n", ops->szRunId);
	fprintf(out, "rdir=%s\n", ops->ldir);
	fprintf(out, "bmap=%s\n", ops->szBmapFileName);
	fprintf(out, "bped=%s\n", ops->szBpedFileName);
	fprintf(out, "phe=%s\n", ops->szPheFileName);
	fprintf(out, "missingness=%s\n", ops->szMissingnessFileName);
	fprintf(out, "founders=%s\n", ops->szFoundersFileName);
	fprintf(out, "exclude=%s\n", ops->szExcludeListFileName);
	fprintf(out, "mpheno=%i\n", ops->mpheno);
	fprintf(out, "pheno-name=%s\n", ops->szPheColumnName);
	fprintf(out, "maf=%f\n", ops->min_af);
	fprintf(out, "geno=%f\n", ops->MAX_GENO_MISSING);
	fprintf(out, "mind=%f\n", ops->MAX_IND_MISSING);
	fprintf(out, "nperm=%i\n", ops->nperm);

	// HWE Test settings
	if (ops->HWD_test) {
		if (!ops->HWD_standard) fprintf(out, "hwe=%i\n", ops->HWD_test);
		else  fprintf(out, "hwe2=%i\n", ops->HWD_test);
		fprintf(out, "hwe_limit=%f\n", ops->HWD_limit);
	}

	if (out != NULL) fclose(out);

	return copyAndDeleteFile(localFilepath, remoteFilepath);
}